/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2022. All rights reserved.
 * Description: 集合通信算子信息库头文件
 * Author: lilianlin
 * Create: 2019-11-28
 */

#ifndef HCOM_OPS_KERNEL_BUILDER_H
#define HCOM_OPS_KERNEL_BUILDER_H

#include "ops_kernel_builder_base.h"
#include "hccl_comm_pub.h" // RankInfo
#include "op_hcom_comm.h"


constexpr u32 ALLTOALLV_INPUT_VEC_SIZE = 4;

namespace hccl {
class HcomOpsKernelBuilder : public HCCLOpsKernelBuilder {
public:
    HcomOpsKernelBuilder();
    ~HcomOpsKernelBuilder() override;
    // memory allocation requirement
    ge::Status CalcOpRunningParam(ge::Node& node) override;
    // generate taskinfo for op
    ge::Status GenerateTask(const ge::Node &node, ge::RunContext &runContext, \
        std::vector<domi::TaskDef> &taskDefList) override;

private:
    HcclResult HcomCalcOpRunningParam(ge::Node& node);
    HcclResult GetSupportedOP(std::vector<std::string> &hcclSupportOp) const override;
    HcclResult GetOpWorkspaceMemSize(ge::Node& node, const std::string &sCollectiveType, u64 &opMemSize);
    HcclResult GetOpScratchMemSize(ge::Node& node, const std::string &sCollectiveType, u64 &opMemSize, u64 count,
        u32 dataTypeSize, HcclDataType dataType, s32 rankSize, int64_t hcomComm, std::string sGroup, s32 serverNum);
    HcclResult SetOpMemAttr(ge::Node& node, const std::string &sCollectiveType, const u64 &opMemSize) override;
    HcclResult SetOpAtomicInputIndex(ge::Node& node, const std::string &sCollectiveType) override;
    HcclResult GetCountFromOpDesc(const ge::OpDescPtr &op, const std::string &sCollectiveType, \
        HcclDataType dataType, u64 &count);
    HcclResult GetHcomReceiveOpOutputSize(const ge::OpDescPtr &op, u32 dataTypeSize, u64 &outputSize);
    HcclResult GetRootGraphID(const ge::Node& node, uint32_t &graphId);
    HcclResult GetCommFromOpDesc(const ge::OpDescPtr &op, int64_t &hcomComm, std::string &sGroup);
    HcclResult GetDestRankFromDesc(const ge::OpDescPtr &op, u32 &destRank, const std::string &sCollectiveType);
    HcclResult GetSrcRankFromDesc(const ge::OpDescPtr &op, u32 &srcRank, const std::string &sCollectiveType);
    HcclResult GetSrTagFromDesc(const ge::OpDescPtr &op, u32 &srTag, const std::string &sCollectiveType);
    HcclResult GetOriginalGraphShapeTypeFromDesc(const ge::OpDescPtr &op, u32 &shapeType);
    HcclResult MarkRemoteAccessMemoryType(ge::Node& node);
    HcclResult GenerateTaskDef(const ge::Node& node, HCCL_KERNEL_INFO_PRIVATE_DEF &privateDefBuf,
        domi::TaskDef &taskDef);
    HcclResult SetAlltoAllVParams(const ge::Node& node, HCCL_ALLTOALLV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf,
        const std::string &sCollectiveType);
    HcclResult SetReduceScatterVParams(const ge::Node& node, HCCL_REDUCESCATTERV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf);
    HcclResult SetAllGatherVParams(const ge::Node& node, HCCL_ALLGATHERV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf);
    HcclResult CopyAlltoAllVParamsToDef(const ge::Node& node, HCCL_ALLTOALLV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf);
    HcclResult CopyReduceScatterVParamsToDef(const ge::Node& node, HCCL_REDUCESCATTERV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf);
    HcclResult CopyAllGatherVParamsToDef(const ge::Node& node, HCCL_ALLGATHERV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf);
    HcclResult SetPrivateDefWithTensorInfo(const ge::Node &node, HCCL_KERNEL_INFO_PRIVATE_DEF &privateDefBuf,
        domi::TaskDef &taskDef);
    HcclResult GetCrackParamsInfo(const ge::Node& node, u32 tensorNum,
        int64_t* tensorOffset, int64_t* tensorSize, int64_t* crackOffset, int64_t* crackSize);
    HcclResult GetTensorParamsInfo(const ge::Node& node, u32 tensorNum,
        int64_t* tensorOffset, int64_t* tensorSize);
    HcclResult SetAlltoAllVDataTypeToDef(const ge::OpDescPtr &op,
        HCCL_ALLTOALLV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf);
    HcclResult CopyAlltoAllVCParamsToDef(const ge::Node& node, HCCL_ALLTOALLV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf);
    HcclResult SetAlltoAllVCDataTypeToDef(const ge::OpDescPtr &op,
        HCCL_ALLTOALLV_KERNEL_INFO_PRIVATE_DEF &privateDefBuf);
    HcclResult GetAlltoAllCountsDispl(const ge::OpDescPtr &op, std::vector<int64_t> &sendCounts,
        std::vector<int64_t> &sendDispls, std::vector<int64_t>& recvCounts, std::vector<int64_t>& recvDispls);
    HcclResult GetReduceScatterVCountsDispl(ge::Node& node, std::vector<int64_t> &sendCounts,
        std::vector<int64_t> &sendDispls, std::vector<int64_t>& recvCount);
    HcclResult GetAllGatherVCountsDispl(ge::Node& node, std::vector<int64_t> &sendCount,
        std::vector<int64_t> &recvCounts, std::vector<int64_t>& recvDispls);
    HcclResult GetAlltoAllCountsDispl(ge::Node& node, std::vector<int64_t> &sendCounts,
        std::vector<int64_t> &sendDispls, std::vector<int64_t>& recvCounts, std::vector<int64_t>& recvDispls);
    HcclResult GetAlltoAllCountsDispl(const ge::OpDescPtr &op, std::vector<int64_t> &sendCountMatrix);
    HcclResult GetAlltoAllCountsDispl(ge::Node& node, std::vector<int64_t> &sendCountMatrix);
    HcclResult GetAlltoAllvStagedScratchMemSize(ge::Node& node, const int64_t &hcomComm, const string &sGroup,
        u32 rankSize, u64 &getMemSize);
    HcclResult GetRanksize(const int64_t &hcomComm, const string &sGroup, s32 &rankSize);
    HcclResult GetLookupUpdateWorkspace(ge::Node& node, u64 &opMemSize, s32 flags);
    bool UseOneLayerAlltoAllv();
    inline HcclResult GetOpIntAttr(const ge::OpDescPtr &op, const string &attr, s32 &output)
    {
        if (ge::AttrUtils::HasAttr(op, attr)) {
            ge::AttrUtils::GetInt(op, attr, output);
        } else {
            return HCCL_E_NOT_FOUND;
        }
        return HCCL_SUCCESS;
    }
    HcclResult GetRemoteLookupAttr(const ge::Node &node, EmbeddingServiceParams &esInfo,
        const std::string &sCollectiveType);
    HcclResult GetNeedMapRankFromDesc(const ge::OpDescPtr &op, bool &needMapRank);
    HcclResult GenerateTaskPrivateDef(const ge::Node &node, HCCL_KERNEL_INFO_PRIVATE_DEF &privateDefBuf,
        domi::TaskDef &taskDef, const std::string sCollectiveType);

    HcclResult GetAndSetTaskNum(const ge::Node &node, const std::string &sCollectiveType,
        const u64 &streamNum);
    void GetUniqueTagFromNode(const ge::Node& node, s32 tag, string &uniqueTag);
    // 从算子中获取计算pipline切分数量的信息
    HcclResult GetOpDataSize(const ge::Node &node, u64 &opDataSize);
    HcclResult CalAndSetOpWorkerSpaceForKnowShape(ge::Node& node, const std::string &sCollectiveType, u64 &opMemSize);
    // 设置附属从流信息
    HcclResult SetAttachedStreamInfoList(ge::Node &node, const std::string &group);
    std::unordered_map<std::string, u32> nodeNameCountMap_;
    std::unordered_map<std::string, u64> nodeTagopMemSizeMap_; // tag, opMemSize
    const std::unordered_set<std::string> ES_OPS = { HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP_PAIRED,
        HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP, HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP_UNIQUED_PAIRED,
        HCCL_KERNEL_OP_TYPE_COLL_REMOTE_UPDATE, HCCL_KERNEL_OP_TYPE_COLL_REMOTE_UPDATE_PAIRED,
        HCCL_KERNEL_OP_TYPE_REMOTE_LOOKUP};
};
}
#endif // GE_OPS_KERNEL_INFO_H
